"""
业务逻辑
"""
import functools
from fastapi import HTTPException


class LoginSafeChecker:
    """登陆安全检查  业务功能"""
    FAIL_KEY_PRE = 'auth:account:fail'
    LOCK_KEY_PRE = 'auth:account:lock'
    LOCK_IP_PRE = 'auth:ip:lock'

    BLACK_IP_KEY = 'auth:black_ips'

    def __init__(self, **kwargs):
        self.lock_window = 60
        self.lock_time = 180
        self.max_retry = 3
        self.is_lock_account = True
        self.is_lock_ip = False
        self.rdb = None

        if kwargs:
            self.init(**kwargs)

    def init(self, lock_window=60, lock_time=180, max_retry=3, is_lock_user=True, is_lock_ip=False, rdb=None):
        self.lock_window = lock_window
        self.lock_time = lock_time
        self.max_retry = max_retry
        self.is_lock_account = is_lock_user
        self.is_lock_ip = is_lock_ip
        self.rdb = rdb

    def is_locked_account(self, account):
        """是否是锁定用户"""
        if self.rdb.redis.get(f'{self.LOCK_KEY_PRE}:{account}'):
            return True
        return False

    def is_locked_ip(self, ip):
        """是否是锁定地址"""
        if self.rdb.redis.get(f'{self.LOCK_IP_PRE}:{ip}'):
            return True
        return False

    def is_black_ip(self, ip):
        """是否是黑名单ip"""
        if self.rdb.redis.sismember(self.BLACK_IP_KEY, ip):
            return True
        return False

    def lock_account(self, account):
        """锁定指定账号"""
        self.rdb.redis.set(f'{self.LOCK_KEY_PRE}:{account}', account, ex=self.lock_time)

    def lock_login_ip(self, ip):
        """锁定登陆IP"""
        self.rdb.redis.set(f'{self.LOCK_IP_PRE}:{ip}', ip, ex=self.lock_time)

    def retry_failed_times(self, account, ip):
        """登陆失败次数计数
        超过指定次数 后锁定
        """
        times = self.rdb.redis.incr(f'{self.FAIL_KEY_PRE}:{account}')
        self.rdb.redis.expire(f'{self.FAIL_KEY_PRE}:{account}', self.lock_window)
        if times >= self.max_retry:
            if self.is_lock_account:
                self.lock_account(account)
            if self.is_lock_ip:
                self.lock_login_ip(ip)
        return times

    def unlock_account(self, account):
        """解锁账号"""
        self.rdb.redis.delete(f'{self.LOCK_KEY_PRE}:{account}')
        self.rdb.redis.delete(f'{self.FAIL_KEY_PRE}:{account}')

    def unlock_ip(self, ip):
        """解锁IP"""
        self.rdb.redis.delete(f'{self.LOCK_IP_PRE}:{ip}')

    def locked_account(self):
        """获取已锁定账号"""
        accounts = [item.decode().replace(f'{self.LOCK_KEY_PRE}:', '')
                    for item in self.rdb.redis.keys(f'{self.LOCK_KEY_PRE}:*')]
        return accounts

    def locked_ip(self):
        """获取已锁定IP"""
        ips = [item.decode().replace(f'{self.LOCK_IP_PRE}:', '')
               for item in self.rdb.redis.keys(f'{self.LOCK_IP_PRE}:*')]
        return ips

    def alert(self, message):
        """登陆告警？ 根据配置"""
        pass

    def check(self, func):
        """登陆环境安全检查  业务逻辑"""
        @functools.wraps(func)
        def deco(*args, **kwargs):
            account = kwargs['account']
            ip = kwargs['ip']
            # ip = kwargs['request'].state.ip

            if self.is_locked_account(account):
                detail = '账号已被锁定'
                self.alert(detail)
                raise HTTPException(400, detail=detail)

            if self.is_locked_ip(ip):
                detail = 'IP已被锁定, 请稍后再试'
                self.alert(detail)
                raise HTTPException(400, detail=detail)

            if self.is_black_ip(ip):
                detail = 'IP已被加入黑名单'
                self.alert(detail)
                raise HTTPException(400, detail=detail)

            try:
                resp = func(*args, **kwargs)
                return resp

            except HTTPException as e:
                self.retry_failed_times(account, ip)
                detail = e.detail
                self.alert(detail)
                raise e

            except Exception as e:
                detail = e
                self.alert(detail)
                raise e

        return deco


lsc = LoginSafeChecker()


if __name__ == '__main__':
    @lsc.check
    def a(**kwargs):
        print('a')
        return 'a'

    from common.package.redis_db import rdb
    rdb.init(host='192.168.101.79', password='MY.io2019')
    lsc.init(rdb=rdb)

    a(account='a', ip='xx')
